{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import torch\n",
    "import random\n",
    "from pydantic import BaseModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PredictionRequest(BaseModel):\n",
    "    \"\"\"\n",
    "    定义预测请求的数据模型。\n",
    "    包含预测所需的所有数据特征，与 Data 对象格式一致。\n",
    "    \"\"\"\n",
    "    x: list                 # 节点特征: [n1, 7], float\n",
    "    edge_index: list       # 边索引: [2, n2], long\n",
    "    edge_attr: list        # 边属性: [n2, 3], float\n",
    "    pos: list              # 节点位置: [n1, 3], float\n",
    "    z: list                # 节点属性: [n1], long\n",
    "    smiles: str            # SMILES 字符串\n",
    "    maccs: list            # MACCS 指纹: [1, 167], long\n",
    "    ecfp: list             # ECFP 指纹: [1, 2048], long\n",
    "    y: int                 # 标签: int\n",
    "    craft_feat: list       # CRAFT 特征: [1, 29], float"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_random_sample(n1: int, n2: int) -> dict:\n",
    "    \"\"\"\n",
    "    生成随机样本数据，符合 PredictionRequest 数据模型。\n",
    "\n",
    "    参数:\n",
    "    n1 (int): 分子原子数（节点数），例如 37\n",
    "    n2 (int): 边数量，例如 76\n",
    "\n",
    "    返回:\n",
    "    dict: 包含 PredictionRequest 所需字段的字典，数据类型和形状符合要求\n",
    "\n",
    "    Raises:\n",
    "    ValueError: 如果 n1 <= 0 或 n2 < 0 或 n2 过大导致无法生成有效边\n",
    "    \"\"\"\n",
    "    # 输入验证\n",
    "    if n1 <= 0:\n",
    "        raise ValueError(\"节点数 n1 必须为正整数\")\n",
    "    if n2 < 0:\n",
    "        raise ValueError(\"边数 n2 必须为非负整数\")\n",
    "    max_edges = n1 * (n1 - 1) // 2  # 无向图的最大边数\n",
    "    if n2 > max_edges:\n",
    "        raise ValueError(f\"边数 n2 ({n2}) 不能超过无向图最大边数 {max_edges}\")\n",
    "\n",
    "    # 1. 节点特征 (x): [n1, 7], torch.float32\n",
    "    x = torch.randn(n1, 7, dtype=torch.float32).tolist()\n",
    "\n",
    "    # 2. 边索引 (edge_index): [2, n2], torch.int64\n",
    "    # 生成有效边，确保边连接的节点索引在 [0, n1-1] 范围内\n",
    "    edge_index = []\n",
    "    if n2 > 0:\n",
    "        # 随机生成 n2 条边，避免自环和重复边\n",
    "        possible_edges = [(i, j) for i in range(n1) for j in range(i + 1, n1)]\n",
    "        random.shuffle(possible_edges)\n",
    "        selected_edges = possible_edges[:min(n2, len(possible_edges))]\n",
    "        edge_index = [[e[0], e[1]] for e in selected_edges] + [[e[1], e[0]] for e in selected_edges]  # 无向图，双向边\n",
    "        edge_index = edge_index[:n2]  # 截取精确的 n2 条边\n",
    "        edge_index = [list(e) for e in zip(*edge_index)]  # 转换为 [2, n2] 格式\n",
    "        if len(edge_index[0]) < n2:  # 如果边数不足，补充随机边\n",
    "            remaining = n2 - len(edge_index[0])\n",
    "            extra_edges = torch.randint(0, n1, (2, remaining), dtype=torch.int64).tolist()\n",
    "            edge_index[0].extend(extra_edges[0])\n",
    "            edge_index[1].extend(extra_edges[1])\n",
    "    else:\n",
    "        edge_index = [[], []]  # 空图情况\n",
    "\n",
    "    # 3. 边属性 (edge_attr): [n2, 3], torch.float32\n",
    "    edge_attr = torch.randn(n2, 3, dtype=torch.float32).tolist()\n",
    "\n",
    "    # 4. 节点位置 (pos): [n1, 3], torch.float32\n",
    "    pos = torch.randn(n1, 3, dtype=torch.float32).tolist()\n",
    "\n",
    "    # 5. 节点属性 (z): [n1], torch.int64\n",
    "    # 假设 z 表示原子序数，随机生成 1-100 之间的整数（常见元素范围）\n",
    "    z = torch.randint(1, 101, (n1,), dtype=torch.int64).tolist()\n",
    "\n",
    "    # 6. SMILES 字符串: str\n",
    "    # 使用简单的占位符，长度随机，实际应用中应使用真实 SMILES\n",
    "    smiles = \"C\" * random.randint(5, 20)\n",
    "\n",
    "    # 7. MACCS 指纹: [1, 167], torch.int64 (二进制)\n",
    "    maccs = torch.randint(0, 2, (1, 167), dtype=torch.float32).tolist()\n",
    "\n",
    "    # 8. ECFP 指纹: [1, 2048], torch.int64 (二进制)\n",
    "    ecfp = torch.randint(0, 2, (1, 2048), dtype=torch.float32).tolist()\n",
    "\n",
    "    # 9. 标签 (y): int\n",
    "    y = random.randint(0, 1)  # 假设二分类问题，标签为 0 或 1\n",
    "\n",
    "    # 10. CRAFT 特征: [1, 29], torch.float32\n",
    "    craft_feat = torch.randn(1, 29, dtype=torch.float32).tolist()\n",
    "\n",
    "    # 返回符合 PredictionRequest 的字典\n",
    "    return {\n",
    "        \"x\": x,\n",
    "        \"edge_index\": edge_index,\n",
    "        \"edge_attr\": edge_attr,\n",
    "        \"pos\": pos,\n",
    "        \"z\": z,\n",
    "        \"smiles\": smiles,\n",
    "        \"maccs\": maccs,\n",
    "        \"ecfp\": ecfp,\n",
    "        \"y\": y,\n",
    "        \"craft_feat\": craft_feat\n",
    "    }\n",
    "\n",
    "# 生成一个随机样本，n1=10（节点数），n2=12（边数）\n",
    "fastapi_data = generate_random_sample(n1=10, n2=12)\n",
    "# fastapi_data\n",
    "# fastapi_data = PredictionRequest(**sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "FastAPI API Response:\n",
      "{'prediction': 7, 'confidence': 1}\n"
     ]
    }
   ],
   "source": [
    "# FastAPI API 请求示例\n",
    "fastapi_url = \"http://localhost:8000/predict\"\n",
    "fastapi_response = requests.post(fastapi_url, json=fastapi_data)\n",
    "fastapi_response\n",
    "print(\"\\nFastAPI API Response:\")\n",
    "print(fastapi_response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gnn_0",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
